import argparse

import numpy as np
import torch
from medpy.filter.binary import largest_connected_component
from skimage.io import imread, imshow, imsave

from dataset import BrainSegmentationDataset as Dataset
from unet import UNet
from utils import (
    gray2rgb,
    outline,
    crop_sample,
    pad_sample,
    resize_sample,
    normalize_volume,
    normalize_mask,
)


def main(args):
    device = torch.device("cpu" if not torch.cuda.is_available() else args.device)

    image = imread(args.image)
    mask = imread(args.mask, as_gray=True)
    image = np.array([image])
    mask = np.array([mask])

    # preprocess
    image, mask = crop_sample((image, mask))
    image, mask = pad_sample((image, mask))
    image, mask = resize_sample((image, mask))

    image = normalize_volume(image)
    mask = normalize_mask(mask)
    mask = mask[..., np.newaxis]

    image = image[0].transpose(2, 0, 1)
    mask = mask[0].transpose(2, 0, 1)

    x = torch.from_numpy(image.astype(np.float32)).unsqueeze(0)
    y_true = mask

    with torch.set_grad_enabled(False):
        unet = UNet(in_channels=Dataset.in_channels, out_channels=Dataset.out_channels)
        state_dict = torch.load(args.weights, map_location=device)
        unet.load_state_dict(state_dict)
        unet.eval()
        unet.to(device)

        y_pred = unet(x.to(device))
        y_pred = y_pred.detach().cpu().numpy()

    # postprocess
    y_pred = np.round(y_pred).astype(int)
    y_pred = largest_connected_component(y_pred)[0]

    o_image = gray2rgb(image[1])  # channel 1 is for FLAIR
    o_image = outline(o_image, y_pred[0], color=[255, 0, 0])
    o_image = outline(o_image, y_true[0], color=[0, 255, 0])
    if args.output:
        imsave(args.output, o_image)
    else:
        imshow(o_image)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Inference for segmentation of brain MRI"
    )
    parser.add_argument("--image", "-i", type=str, help="Input Image")
    parser.add_argument("--mask", "-m", type=str, help="Ground Truth Mask")
    parser.add_argument("--output", "-o", type=str, default=None, help="Output file")
    parser.add_argument(
        "--device",
        type=str,
        default="cuda:0",
        help="device for training (default: cuda:0)",
    )
    parser.add_argument(
        "--weights", type=str, required=True, help="path to weights file"
    )
    parser.add_argument(
        "--image-size",
        type=int,
        default=256,
        help="target input image size (default: 256)",
    )
    args = parser.parse_args()
    main(args)
